// BSD 3- Clause License Copyright (c) 2024, Tecorigin Co., Ltd. All rights
// reserved.
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
// Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
// Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation
// and/or other materials provided with the distribution.
// Neither the name of the copyright holder nor the names of its contributors
// may be used to endorse or promote products derived from this software
// without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)  ARISING IN ANY
// WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
// OF SUCH DAMAGE.

#include "interface/include/builtin_type.h"
#include "interface/common/check.h"
#include "interface/common/marco.h"
#include "ual/ops/conv_forward/conv_forward.hpp"

using tecoal::ual::args::ConvFwdArgs;
using tecoal::ual::args::ConvFwdPatchArgs;
using tecoal::ual::ops::ConvFwdOp;
using namespace tecoal;

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define ALIGN_UNIT(in, unit) (((in) + (unit)-1) / (unit) * (unit))
#define SEG_NUM(in, unit) (((in) + (unit)-1) / (unit))
#define BACK_INFER_H(in) (((in)-1) * (SH) + (R))
#define BACK_INFER_W(in) (((in)-1) * (SW) + (S))

#define DB_MAX_USED_SPM_SIZE 239616  // 234KB for SPM

static tecoalStatus_t getConvFwdArgs(tecoalHandle_t handle, const tecoalTensorDescriptor_t xDesc,
                                     const tecoalFilterDescriptor_t wDesc,
                                     const tecoalConvolutionDescriptor_t convDesc,
                                     const tecoalTensorDescriptor_t yDesc, ConvFwdArgs *arg,
                                     ConvFwdPatchArgs *args_patch) {
    arg->spa_num = handle->spa_num;
    arg->spe_num = handle->spe_num;
    arg->N = xDesc->n, arg->C = xDesc->c, arg->H = xDesc->h, arg->W = xDesc->w;
    arg->M = wDesc->m, arg->R = wDesc->r, arg->S = wDesc->s;
    arg->E = yDesc->h, arg->F = yDesc->w;
    arg->pad_h = convDesc->padA[0];
    arg->pad_w = convDesc->padA[1];
    arg->stride_h = convDesc->filterStrideA[0];
    arg->stride_w = convDesc->filterStrideA[1];
    arg->dilation_h = convDesc->dilationA[0];
    arg->dilation_w = convDesc->dilationA[1];
    arg->alpha = 1.0;
    arg->beta = 0.0;

    arg->out_data_type = Convert::toUALDataType(yDesc->dataType);

    arg->x = nullptr;
    arg->w = nullptr;
    arg->y = nullptr;

    arg->workSpace = nullptr;
    arg->workSpaceSize = 0;

    args_patch->convf = arg;
    args_patch->x_data_type = Convert::toUALDataType(xDesc->dataType);
    args_patch->w_data_type = Convert::toUALDataType(wDesc->dataType);
    args_patch->y_data_type = Convert::toUALDataType(yDesc->dataType);

    return TECOAL_STATUS_SUCCESS;
}

// Create an instance of convolution descriptor
tecoalStatus_t TECOALWINAPI
tecoalCreateConvolutionDescriptor(tecoalConvolutionDescriptor_t *convDesc) {
    *convDesc = (tecoalConvolutionDescriptor_t)malloc(sizeof(tecoalConvolutionStruct));
    (*convDesc)->padA = nullptr;
    (*convDesc)->filterStrideA = nullptr;
    (*convDesc)->dilationA = nullptr;
    (*convDesc)->mathType = TECOAL_TENSOR_ACC_MATH;
    return TECOAL_STATUS_SUCCESS;
}

// Destroy an instance of convolution descriptor
tecoalStatus_t TECOALWINAPI
tecoalDestroyConvolutionDescriptor(tecoalConvolutionDescriptor_t convDesc) {
    if (convDesc != NULL) {
        free(convDesc);
    }
    return TECOAL_STATUS_SUCCESS;
}

// Set the parameters for a 2D convolution operation in its descriptor.
tecoalStatus_t TECOALWINAPI tecoalSetConvolution2dDescriptor(
    tecoalConvolutionDescriptor_t convDesc, int pad_h,  // zero-padding height
    int pad_w,                                          // zero-padding width
    int u,                                              // vertical filter stride
    int v,                                              // horizontal filter stride
    int dilation_h,                                     // filter dilation in the vertical dimension
    int dilation_w,  // filter dilation in the horizontal dimension
    tecoalConvolutionMode_t mode, tecoalDataType_t dataType) {
    convDesc->pad_h = pad_h;
    convDesc->pad_w = pad_w;
    convDesc->u = u;
    convDesc->v = v;
    convDesc->dilation_h = dilation_h;
    convDesc->dilation_w = dilation_w;
    if (convDesc->padA != nullptr) {
        free(convDesc->padA);
    }
    convDesc->padA = (int *)malloc(sizeof(int) * 2);
    convDesc->padA[0] = pad_h;
    convDesc->padA[1] = pad_w;
    if (convDesc->filterStrideA != nullptr) {
        free(convDesc->filterStrideA);
    }
    convDesc->filterStrideA = (int *)malloc(sizeof(int) * 2);
    convDesc->filterStrideA[0] = u;
    convDesc->filterStrideA[1] = v;
    if (convDesc->dilationA != nullptr) {
        free(convDesc->dilationA);
    }
    convDesc->dilationA = (int *)malloc(sizeof(int) * 2);
    convDesc->dilationA[0] = dilation_h;
    convDesc->dilationA[1] = dilation_w;
    convDesc->mode = mode;
    convDesc->dataType = dataType;
    return TECOAL_STATUS_SUCCESS;
}

// Retrieve the parameters of a 2D convolution descriptor.
tecoalStatus_t TECOALWINAPI tecoalGetConvolution2dDescriptor(
    const tecoalConvolutionDescriptor_t convDesc, int *pad_h,  // zero-padding height
    int *pad_w,                                                // zero-padding width
    int *u,                                                    // vertical filter stride
    int *v,                                                    // horizontal filter stride
    int *dilation_h,  // filter dilation in the vertical dimension
    int *dilation_w,  // filter dilation in the horizontal dimension
    tecoalConvolutionMode_t *mode, tecoalDataType_t *dataType) {
    if (pad_h != NULL) *pad_h = convDesc->pad_h;
    if (pad_w != NULL) *pad_w = convDesc->pad_w;
    if (u != NULL) *u = convDesc->u;
    if (v != NULL) *v = convDesc->v;
    if (dilation_h != NULL) *dilation_h = convDesc->dilation_h;
    if (dilation_w != NULL) *dilation_w = convDesc->dilation_w;
    if (mode != NULL) *mode = convDesc->mode;
    if (dataType != NULL) *dataType = convDesc->dataType;
    return TECOAL_STATUS_SUCCESS;
}

// Calculate the size of the workspace needed for a forward convolution operation.
tecoalStatus_t TECOALWINAPI tecoalGetConvolutionForwardWorkspaceSize(
    tecoalHandle_t handle, const tecoalTensorDescriptor_t xDesc,
    const tecoalFilterDescriptor_t wDesc, const tecoalConvolutionDescriptor_t convDesc,
    const tecoalTensorDescriptor_t yDesc, tecoalAlgo_t algo, size_t *workSpaceSizeInBytes) {
    ConvFwdArgs arg;
    ConvFwdPatchArgs args_patch;
    checkTecoalStatus(getConvFwdArgs(handle, xDesc, wDesc, convDesc, yDesc, &arg, &args_patch));
    args_patch.algo = Convert::toUalAlgoType(algo);
    ConvFwdOp op{};
    Status status = op.getWorkspace(&args_patch, workSpaceSizeInBytes);
    checkUalStatusInTecoal(status);
    return TECOAL_STATUS_SUCCESS;
}

static int calcuBlockParam(const ConvFwdArgs &arg, int *block_c_p, int *block_e_p,
                                      int *block_f_p) {
    const int N = arg.N, C = arg.C, H = arg.H, W = arg.W, M = arg.M, R = arg.R, S = arg.S,
              E = arg.E, F = arg.F, SH = arg.stride_h, SW = arg.stride_w;
    int &block_c = *block_c_p;
    int &block_e = *block_e_p;
    int &block_f = *block_f_p;
    // All splitting strategies are designed based on the 234KB SPM
    constexpr int unit_block_c = 32, unit_block_m = 32;
    const int min_block_f = MIN(F, 32), max_block_f = MIN(F, 128);
    const int max_block_e = E;
    int start_block_e = 1;
    int start_block_c = ALIGN_UNIT(C, unit_block_c);
    int all_c = ALIGN_UNIT(C, unit_block_c);
    // Keep bf even if C is odd and half is half
    int enlarge_bf = sizeof(half) == 2 && C % 2 ? 2 : 1;
    int start_block_f = (F <= 96) ? max_block_f : min_block_f;
    int start_block_f0 = start_block_f;
    // 1. First fix bF and bE, and check if it is possible to split with bC = C
    while (sizeof(half) * BACK_INFER_H(start_block_e) * BACK_INFER_W(start_block_f0) * all_c +
                   sizeof(half) * unit_block_m * all_c * 2 +
                   sizeof(half) * start_block_e * start_block_f0 * unit_block_m * 2 >
               DB_MAX_USED_SPM_SIZE &&
           (start_block_f0 + 1) / 2 >= 16) {
        start_block_f0 = (start_block_f0 + 1) / 2;
    }
    int read_all_C =
        sizeof(half) * BACK_INFER_H(start_block_e) * BACK_INFER_W(start_block_f0) * all_c +
        sizeof(half) * unit_block_m * all_c * 2 +
        sizeof(half) * start_block_e * start_block_f0 * unit_block_m * 2;
    // 2. If 1 is possible, then further check if bE and bF can be enlarged
    if (read_all_C < DB_MAX_USED_SPM_SIZE) {
        block_f = start_block_f0;
        block_c = start_block_c;
        block_e = start_block_e;
        // First enlarge max_be
        for (int tmp_be = start_block_e; tmp_be <= max_block_e; tmp_be += 1) {
            read_all_C = sizeof(half) * BACK_INFER_H(tmp_be) * BACK_INFER_W(block_f) * all_c +
                         sizeof(half) * unit_block_m * all_c * 2 +
                         sizeof(half) * tmp_be * block_f * unit_block_m * 2;
            if (read_all_C > DB_MAX_USED_SPM_SIZE || tmp_be * block_f > 128) break;
            block_e = tmp_be;  // enlarge be
        }
        // Then enlarge tmp_bf
        for (int tmp_bf = start_block_f0; tmp_bf <= max_block_f; tmp_bf += enlarge_bf) {
            read_all_C = sizeof(half) * BACK_INFER_H(block_e) * BACK_INFER_W(tmp_bf) * all_c +
                         sizeof(half) * unit_block_m * all_c * 2 +
                         sizeof(half) * block_e * tmp_bf * unit_block_m * 2;
            if (read_all_C > DB_MAX_USED_SPM_SIZE || block_e * tmp_bf > 128) break;
            block_f = tmp_bf;  // enlarge bf
        }
        block_e = MIN(block_e, SEG_NUM(E, SEG_NUM(E, block_e)));
        block_f = MIN(block_f, SEG_NUM(F, SEG_NUM(F, block_f)));

        int new_block_e = MIN(SEG_NUM(ALIGN_UNIT(E, 4), 4), block_e);
        while (N * SEG_NUM(E, block_e) * SEG_NUM(F, block_f) <= 16) {
            block_e = (block_e + 1) / 2;
            if (block_e == 1) break;
        }
        int loops = N * SEG_NUM(E, block_e) * SEG_NUM(F, block_f);
        float eff1 = (float)loops / (float)ALIGN_UNIT(loops, 32);
        loops = N * SEG_NUM(E, new_block_e) * SEG_NUM(F, block_f);
        float eff2 =
            float(SEG_NUM(M, 32) * loops) / float(SEG_NUM(M, 256) * SEG_NUM(loops, 4) * 32);
        if (block_e * block_f < 32 && block_e != new_block_e && eff1 < eff2) {
            block_e = new_block_e;
        }
        return 1;  // reading all C
    }
    // 3. If 1 is not possible, then fix bC and adjust bE and bF for splitting
    // Split C
    int max_block_c = unit_block_c;
    block_c = max_block_c;    // setting bc
    block_e = start_block_e;  // setting be
    block_f = start_block_f;  // setting bf

    // First enlarge be
    for (int tmp_be = start_block_e; tmp_be <= max_block_e; tmp_be += 1) {
        read_all_C = sizeof(half) * BACK_INFER_H(tmp_be) * BACK_INFER_W(block_f) * max_block_c +
                     sizeof(half) * unit_block_m * max_block_c * 2 +
                     sizeof(half) * tmp_be * block_f * unit_block_m * 2;
        if (read_all_C > DB_MAX_USED_SPM_SIZE || tmp_be * block_f > 128) break;
        block_e = tmp_be;
    }

    // Then enlarge tmp_bf
    for (int tmp_bf = start_block_f; tmp_bf <= max_block_f; tmp_bf += enlarge_bf) {
        read_all_C = sizeof(half) * BACK_INFER_H(block_e) * BACK_INFER_W(tmp_bf) * max_block_c +
                     sizeof(half) * unit_block_m * max_block_c * 2 +
                     sizeof(half) * block_e * tmp_bf * unit_block_m * 2;
        if (read_all_C > DB_MAX_USED_SPM_SIZE || block_e * tmp_bf > 128) break;
        block_f = tmp_bf;
    }
    block_e = MIN(block_e, SEG_NUM(E, SEG_NUM(E, block_e)));
    block_f = MIN(block_f, SEG_NUM(F, SEG_NUM(F, block_f)));

    if (SEG_NUM(E, block_e) * SEG_NUM(F, block_f) <= 16 && (block_f + 1) / 2 >= 16) {
        block_f = (block_f + 1) / 2;
    } else if (SEG_NUM(E, block_e) * SEG_NUM(F, block_f) <= 16) {
        block_e = (block_e + 1) / 2;
    }
    // Lastly, enlarge bc
    while (sizeof(half) * BACK_INFER_H(block_e) * BACK_INFER_W(block_f) * max_block_c +
               sizeof(half) * unit_block_m * max_block_c * 2 +
               sizeof(half) * block_e * block_f * unit_block_m * 2 <=
           DB_MAX_USED_SPM_SIZE) {
        max_block_c += unit_block_c;
    }
    max_block_c -= unit_block_c;
    max_block_c = MIN(max_block_c, ALIGN_UNIT(SEG_NUM(C, SEG_NUM(C, max_block_c)), 32));

    // Ensure last BC is less than 32
    if (C % 32 == 0) {
        max_block_c = MIN(max_block_c, C);
    } else {
        max_block_c = MIN(max_block_c, C - C % 32);
        max_block_c = MAX(32, max_block_c);
    }
    block_c = max_block_c;  // setting bc
    return 1;
}

// Function to perform the forward pass for batch convolution
tecoalStatus_t TECOALWINAPI tecoalConvolutionForward(
    tecoalHandle_t handle, const void *alpha, const tecoalTensorDescriptor_t xDesc, const void *x,
    const tecoalFilterDescriptor_t wDesc, const void *w,
    const tecoalConvolutionDescriptor_t convDesc, tecoalAlgo_t algo, void *workSpace,
    size_t workSpaceSizeInBytes, const void *beta, const tecoalTensorDescriptor_t yDesc, void *y) {
    // Structure to hold arguments for the forward convolution operation
    ConvFwdArgs arg;

    // Structure to hold patched arguments for the operation
    ConvFwdPatchArgs args_patch;

    // Validate and retrieve the convolution operation arguments based on descriptors
    checkTecoalStatus(getConvFwdArgs(handle, xDesc, wDesc, convDesc, yDesc, &arg, &args_patch));
    int block_c = 0, block_e = 0, block_f = 0;
    if (!calcuBlockParam(arg, &block_c, &block_e, &block_f)) {
        return TECOAL_STATUS_SUCCESS;
    }

    // Assign data to arguments
    //  printf("block_c = %d, block_e = %d, block_f = %d\n", block_c, block_e, block_f);
    arg.block_c = block_c;
    arg.block_e = block_e;
    arg.block_f = block_f;
    arg.x = x;
    arg.w = w;
    arg.y = y;
    arg.workSpace = workSpace;
    arg.workSpaceSize = workSpaceSizeInBytes;
    args_patch.algo = Convert::toUalAlgoType(algo);

    // Execute the forward convolution operation
    RUN_OP(ConvFwdOp, arg, args_patch, handle);

    return TECOAL_STATUS_SUCCESS;
}
